import copy
import json

import matplotlib.pyplot as plt
import networkx as nx


class GraphFilter():

    def write2json(self, filename='./data/MC1_NEW_filtered.json'):
        dict_to_write = {
            "nodes": [self.dict_name_to_node[n] for n in self.point_list],
            "links": [self.dict_name_to_edge[l] for l in self.link_list]
        }
        with open(filename, 'w') as write_f:
            write_f.write(json.dumps(
                dict_to_write, indent=4, ensure_ascii=True))

    def filter_non_connected_subgraph(self):
        point_list = list(self.dict_name_to_node.keys())
        link_list = list(self.dict_name_to_edge.keys())

        G = nx.Graph()
        # 转化为图结构
        for node in point_list:
            G.add_node(node)

        for link in link_list:
            G.add_edge(link[0], link[1])

        for c in nx.connected_components(G):
            nodeSet = G.subgraph(c).nodes()
            # print(len(nodeSet))  # 打印每个连通子图的节点数量
            # 删除其中的点和边
            if (len(nodeSet) < 100):
                for del_node in nodeSet:
                    self.point_list.remove(del_node)
        self.link_list = []
        for (n1, n2) in link_list:
            if n1 in self.point_list and n2 in self.point_list:
                self.link_list.append((n1, n2))
        # print(len(self.link_list))

    def filter_low_degree_nodes(self, filter_degree=2):
        # filter_degree融合点的最小度
        point_list = copy.deepcopy(self.point_list)
        link_list = copy.deepcopy(self.link_list)

        G = nx.Graph()
        # 转化为图结构
        for node in point_list:
            G.add_node(node)

        for link in link_list:
            G.add_edge(link[0], link[1])

        for id, degree in G.degree:
            if degree == 1:
                self.point_list.remove(id)
            elif degree <= filter_degree and self.dict_name_to_node[id]["type"] == "person":
                self.point_list.remove(id)
                # 都与该person相连的其他节点连边
                neighbors = [n for n in G.neighbors(id)]
                for nb1 in neighbors:
                    for nb2 in neighbors:
                        if nb1 == nb2:
                            continue
                        info1 = self.dict_name_to_edge.get(
                            (id, nb1), self.dict_name_to_edge.get((nb1, id), None))
                        info2 = self.dict_name_to_edge.get(
                            (id, nb2), self.dict_name_to_edge.get((nb2, id),None))
                        if info1 is None or info2 is None:
                            continue
                        # 认为是一种partnership
                        if info1["source"] == info2["target"]:
                            self.link_list.append((nb2, nb1))
                            self.dict_name_to_edge[(nb2, nb1)] = {
                                "type": "partnership",
                                "source": nb2,
                                "target": nb1,
                                "key": info2["key"],
                                "value": info1["value"]*info2["value"]
                            }
                        elif info2["source"] == info1["target"]:
                            self.link_list.append((nb1, nb2))
                            self.dict_name_to_edge[(nb1, nb2)] = {
                                "type": "partnership",
                                "source": nb1,
                                "target": nb2,
                                "key": info1["key"],
                                "value": info1["value"]*info2["value"]
                            }
                        else:
                            self.link_list.append((nb2, nb1))
                            self.dict_name_to_edge[(nb2, nb1)] = {
                                "type": "partnership",
                                "source": nb2,
                                "target": nb1,
                                "key": info2["key"],
                                "value": info1["value"]*info2["value"]
                            }
                            self.link_list.append((nb1, nb2))
                            self.dict_name_to_edge[(nb1, nb2)] = {
                                "type": "partnership",
                                "source": nb1,
                                "target": nb2,
                                "key": info1["key"],
                                "value": info1["value"]*info2["value"]
                            }
                            
        self.link_list = []
        for (n1, n2) in link_list:
            if n1 in self.point_list and n2 in self.point_list:
                self.link_list.append((n1, n2))
        # print(len(self.link_list))

    def filter(self, data, filter_degree=2):
        dict_name_to_node = {}
        for n in data['nodes']:
            dict_name_to_node[n['id']] = n

        dict_name_to_edge = {}
        for l in data['links']:
            dict_name_to_edge[(l['source'], l['target'])] = l
        self.dict_name_to_node = dict_name_to_node
        self.dict_name_to_edge = dict_name_to_edge
        self.point_list = list(self.dict_name_to_node.keys())
        self.link_list = list(self.dict_name_to_edge.keys())

        self.filter_non_connected_subgraph()
        # self.write2json(
        #     './data/MC1_NEW_filtered_small_nonconnected_subgraph.json')
        self.filter_low_degree_nodes(filter_degree=filter_degree)
        # self.write2json(
        #     './data/MC1_NEW_filtered_small_nonconnected_subgraph_filter_degree2.json')
        return {
            "nodes": [self.dict_name_to_node[n] for n in self.point_list],
            "links": [self.dict_name_to_edge[l] for l in self.link_list]
        }


if __name__ == "__main__":
    graphfilter = GraphFilter()
    data = graphfilter.filter(
        json.load(open('./data/graph.json', 'r', encoding='utf-8')))
    print(len(data["nodes"]))
    print(len(data["links"]))
